10654. Unique color

 

Given a tree with n vertices numbered 1 through n. The i-th edge connects Vertex ai and Vertex bi. Vertex i is painted in color ci (in this problem, colors are represented as integers).

Vertex x is said to be good when the shortest path from Vertex 1 to Vertex x does not contain a vertex painted in the same color as Vertex x, except for Vertex x itself.

Find all the good vertices.

 

Input. The first line contains the number of vertices n (2 ≤ n ≤ 105). The second line contains colors c1, c2, ..., cn (1 ≤ ci ≤ 105). Each of the next n – 1 lines contains two integers ai and bi (1 ≤ ai, bi ≤ n).

 

Output. Print all good vertices as integers, in ascending order. Print each number on a separate line.

 

Sample input 1

Sample output 1

6

2 7 1 8 2 8

1 2

3 6

3 2

4 3

2 5

1

2

3

4

6

 

 

Sample input 2

Sample output 2

10

3 1 4 1 5 9 2 6 5 3

1 2

2 3

3 4

4 5

5 6

6 7

7 8

8 9

9 10

1

2

3

5

6

7

8

 

 

SOLUTION

graphs – depth first search

 

Algorithm analysis

Start the depth-first search from vertex 1. When entering vertex v of color[v], we increase the value of used[color[v]] by 1. The value used[color[v]] contains the number of times that the vertex of color[v] has met on the path from 1 to v (including the vertex v itself). If the color[v] is encountered only once on the path, then vertex v is good, add it to the final set.

When leaving the vertex v, the value used[color[v]] should be decreased by 1.

 

Example

The graph from the first sample looks like this:

The vertex 5 is not good because on the path 1 – 2 – 5 vertices 5 and 1 have the same color.

The vertex 6 is good because on the path 1 – 2 – 3 – 6 the vertex colors are different from the color of the vertex 6.

 

Algorithm realization

Store the input graph in the adjacency list g. Declare the arrays.

 

vector<int> used, color;

vector<vector<int>> g;

set<int> st;

 

The dfs function implements a depth-first search. The variable par is the ancestor of v.

 

void dfs(int v, int par)

{

 

Vertex v has color color[v]. Note that on the way from vertex 1, we met a vertex of color[v].

 

  used[color[v]]++;

 

The value of used[color[v]] contains the number of times that the vertex of color[v] met on the path from 1 to v (including the vertex v itself). If color[v] is encountered only once on the path, then the vertex v is good, store it in the result set st.

 

  if (used[color[v]] == 1) st.insert(v);

 

Iterate over all the edges (v, to) outgoing from v. If to is not an ancestor of v (to ≠ par) then run the depth-first search from to. In this case, the ancestor of to becomes v.

 

  for (int to : g[v])

    if (to != par) dfs(to, v);

 

When leaving the vertex v, decrease the value of used[color[v]] by 1.

 

  used[color[v]]--;

}

 

The main part of the program. Read the number of vertices n and an array of colors.

 

scanf("%d", &n);

color.resize(n + 1);

for (i = 1; i <= n; i++)

  scanf("%d", &color[i]);

 

Read the graph.

 

used.resize(100001);

g.resize(n + 1);

for (i = 1; i < n; i++)

{

  scanf("%d %d", &x, &y);

  g[x].push_back(y);

  g[y].push_back(x);

}

 

Run the depth-first search from vertex 1.

 

dfs(1, 1);

 

Print the good vertices.

 

for (int val : st)

  printf("%d\n", val);